import * as tf from "@tensorflow/tfjs";
import "@tensorflow/tfjs-backend-webgpu";
import * as tfvis from "@tensorflow/tfjs-vis";
import { baseURL } from "@/renderer/api/machineLearningAPI";
import {
  NUM_CLASSES,
  CLASSES,
  IMAGE_WIDTH,
  IMAGE_HEIGHT,
} from "../datasets/OCRDatasetConstants";
import OCRData from "./OCRData";

const moduleName = "ocr";
const loadURL = `${baseURL}/models/ocr/model.json`;
const saveURL = `${baseURL}/savemodel/${moduleName}`;

tf.setBackend("webgpu").then((res) => {
  // console.log('tfjs-backend-webgpu');
});

export default class OCRModel {
  model;
  BATCH_SIZE = 32; // 32-512 单周期单批数
  EPOCHS = 10; // 周期数

  constructor() {
    this.initVis();
    tfvis.visor().close();
  }

  initVis() {
    tfvis.visor();

    const visor: HTMLElement = document.querySelector(
      "#tfjs-visor-container > div"
    );
    visor.style.setProperty("height", "calc(100% - 30px)");
    visor.style.setProperty("top", "auto");
    visor.style.setProperty("bottom", "0");
  }

  async load() {
    try {
      const model = await tf.loadLayersModel(loadURL);
      this.model = model;
    } catch (e) {
      this.model = await createOCRModel();
    }

    this.compile();
  }

  async save() {
    await this.model.save(
      tf.io.http(saveURL, {
        requestInit: { method: "POST" },
      })
    );
  }
  // fit之前必须compile
  compile() {
    const LEARNING_RATE = 0.15;
    const optimizer = tf.train.sgd(LEARNING_RATE);
    // 独热表示法categoricalCrossentropy, 稀疏表示法（索引表示法）sparseCategoricalCrossentropy
    this.model.compile({
      optimizer: optimizer,
      loss: "sparseCategoricalCrossentropy",
      metrics: ["accuracy"],
    });
  }

  onBatchEnd(batch, logs) {}
  onEpochEnd(epoch, logs) {}

  // 训练
  async train(data, watch?: boolean) {
    tfvis.visor().open();

    const SHUFFLE = true;

    data.trainIndex = 0;
    data.testIndex = 0;

    const callbacks = [];

    callbacks.push({
      onBatchEnd: this.onBatchEnd,
      onEpochEnd: this.onEpochEnd,
    });

    if (watch) callbacks.push(fitCallbacks());
    // model.fit函数
    // const numElements = Math.min(3600, data.numTrainElements);

    // const [trainXs, trainYs] = tf.tidy(() => {
    //   const d = data.getTrainData(numElements);
    //   return [d.xs, d.labels];
    // });

    // const [testXs, testYs] = tf.tidy(() => {
    //   const d = data.getTestData(
    //     numElements * (data.numTestElements / data.numTrainElements)
    //   );
    //   return [d.xs, d.labels];
    // });

    // const trainingHistory = await this.model.fit(trainXs, trainYs, {
    //   batchSize: this.BATCH_SIZE,
    //   validationData: [testXs, testYs],
    //   epochs: this.EPOCHS,
    //   shuffle: SHUFFLE,
    //   callbacks,
    // });

    // model.fitDataset函数
    const [trainDataset, testDataset] = tf.tidy(() => {
      return [data.getTrainDataset(this.BATCH_SIZE), data.getTestDataset(this.BATCH_SIZE)]
    })

    const trainingHistory = await this.model.fitDataset(trainDataset, {
      validationData: testDataset,
      epochs: this.EPOCHS,
      callbacks,
    });

    if (!watch) showTrainingHistory(trainingHistory);

    await this.save();
  }
  // 评估
  async evaluate(data) {
    tfvis.visor().open();

    data.testIndex = 0;

    const numElements = Math.min(5000, data.numTestElements);

    const testData = data.getTestData(numElements);
    const testxs = testData.xs;

    // 独热表示法标签转索引argMax
    // const labels = testData.labels.argMax([-1]);
    // 稀疏表示法标签
    const labels = testData.labels;
    const preds = this.model.predict(testxs).argMax([-1]);

    testxs.dispose();

    await showAccuracy([preds, labels], CLASSES);
  }
  // 预测
  async predict(canvas) {
    const ctx = canvas.getContext("2d");
    const data = ctx.getImageData(0, 0, IMAGE_WIDTH, IMAGE_HEIGHT).data;
    const xs = OCRData.imagesDataHandle(data, 1);
    const labelIndex = this.model.predict(xs).argMax([-1]);
    const label = CLASSES.at((await labelIndex.array())[0]);

    return label;
  }
}

function fitCallbacks() {
  const metrics = ["loss", "val_loss", "acc", "val_acc"];
  const container = {
    name: "监控训练",
    tab: "训练",
    styles: { height: "1000px" },
  };
  const callbacks = tfvis.show.fitCallbacks(container, metrics);

  return callbacks;
}

function showTrainingHistory(trainingHistory) {
  tfvis.show.history({ name: "训练历史", tab: "训练" }, trainingHistory, [
    "loss",
    "val_loss",
    "acc",
    "val_acc",
  ]);
}

async function showAccuracy(prediction, classNames) {
  const [preds, labels] = prediction;
  const classAccuracy = await tfvis.metrics.perClassAccuracy(labels, preds);
  const container = { name: "准确率", tab: "评估" };
  tfvis.show.perClassAccuracy(container, classAccuracy, classNames);

  labels.dispose();
}

// 创建并保存模型
export async function createOCRModel() {
  const model = tf.sequential();
  // 在卷积神经网络的第一层中指定输入形状。
  // 然后为指定一些参数在该层中发生的卷积运算。
  model.add(
    tf.layers.conv2d({
      inputShape: [28, 28, 1],
      kernelSize: 5,
      filters: 8,
      strides: 1,
      activation: "relu",
      kernelInitializer: "varianceScaling",
    })
  );

  model.add(tf.layers.maxPooling2d({ poolSize: [2, 2], strides: [2, 2] }));
  model.add(
    tf.layers.conv2d({
      kernelSize: 5,
      filters: 16,
      strides: 1,
      activation: "relu",
      kernelInitializer: "varianceScaling",
    })
  );

  model.add(tf.layers.maxPooling2d({ poolSize: [2, 2], strides: [2, 2] }));
  // 将2D滤波器的输出平坦化为1D矢量
  // 以准备它用于输入到我们的最后一层
  model.add(tf.layers.flatten());
  // 最后一层是一个密集层，它有NUM_CLASSES个输出单元
  model.add(
    tf.layers.dense({
      units: NUM_CLASSES,
      kernelInitializer: "varianceScaling",
      activation: "softmax",
    })
  );

  await model.save(
    tf.io.http(saveURL, {
      requestInit: { method: "POST" },
    })
  );

  return model;
}
